import logging

class HuffmanTreeNode:
    def __init__(self, iIsLeaf = False, iValue = None):
        self.mChildren = [None, None]
        
        self.mIsLeaf = iIsLeaf
        self.mValue = iValue
        
class HuffmanTree:
    def __init__(self):
        self.mRootNode = HuffmanTreeNode()
    
    # Max = 31 bits
    def registerCode(self, iBitString, iValue):
        lCurrentNode = self.mRootNode
        for lBit in iBitString:
            lBitInt = int(lBit)
            if not(lCurrentNode.mChildren[lBitInt]):
                lCurrentNode.mChildren[lBitInt] = HuffmanTreeNode()
            elif lCurrentNode.mChildren[lBitInt].mIsLeaf:
                logging.error("Reached a leaf while trying to insert the code: " + iBitString + " - value: " + str(iValue))
                # Error, reached a Leaf while trying to insert an element!
                return False
            lCurrentNode = lCurrentNode.mChildren[lBitInt]
        
        for lIndex in range(2):
            if lCurrentNode.mChildren[lIndex]:
                logging.error("Trying to add a leaf where we already have an active node for code: " + iBitString + " - value: " + str(iValue))
                # Error, trying to set a leaf on a node that is not a leaf!
                return False
        logging.debug("Successfully added code: " + iBitString + " - value: " + str(iValue))
        lCurrentNode.mIsLeaf = True
        lCurrentNode.mValue = iValue
        return True
    
    # Max = 31 bits
    def readCode(self, iBitString, iPos):
        lCurrentNode = self.mRootNode
        lNumberOfBitsUsed = 0
        for lIndex in range(iPos, len(iBitString)):
            lBitInt = int(iBitString[lIndex])
            lNumberOfBitsUsed += 1
            if not(lCurrentNode.mChildren[lBitInt]):
                logging.error("Reached a dead end while looking for code starting with: " + iBitString[iPos:iPos + lNumberOfBitsUsed])
                # Reached dead end while searching for code!
                return None, lNumberOfBitsUsed
            elif lCurrentNode.mChildren[lBitInt].mIsLeaf:
                logging.debug("Found matching code: " + iBitString[iPos:iPos + lNumberOfBitsUsed] + " - value: " + str(lCurrentNode.mChildren[lBitInt].mValue))
                # We found the code
                return lCurrentNode.mChildren[lBitInt].mValue, lNumberOfBitsUsed
            lCurrentNode = lCurrentNode.mChildren[lBitInt]
        # After all the bits we didn't find a match!
        logging.error("Reached the end of the tree without finding the code starting with: " + iBitString[iPos:iPos + lNumberOfBitsUsed])
        return None, lNumberOfBitsUsed

def createHuffmanTreeFromSymbolDict(iSymbolDict): # No more than 31 bits code
    lHuffmanTree = HuffmanTree()
    lCode = 0
    for lIndex in range(32):
        if lIndex in iSymbolDict:
            for lSymbol in reversed(iSymbolDict[lIndex]):
                lHuffmanTree.registerCode(bin(lCode).lstrip('0b').rjust(lIndex, '0'), lSymbol)
                lCode -= 1
        lCode = (lCode << 1) + 1
    return lHuffmanTree
        